package com.example.socks5_speed;

import android.net.VpnService;

import org.pcap4j.packet.DnsDomainName;
import org.pcap4j.packet.DnsPacket;
import org.pcap4j.packet.DnsQuestion;
import org.pcap4j.packet.DnsRDataA;
import org.pcap4j.packet.DnsResourceRecord;
import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.LengthBuilder;
import org.pcap4j.packet.UdpPacket;
import org.pcap4j.packet.namednumber.DnsClass;
import org.pcap4j.packet.namednumber.DnsOpCode;
import org.pcap4j.packet.namednumber.DnsRCode;
import org.pcap4j.packet.namednumber.DnsResourceRecordType;

import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

public class UdpConnect implements PacketTool.NetworkChannel, BytesParse.Parse {

    private final BytesParse bytesParse;
    private final DatagramSocket datagramSocket;
    private final IpPacket dstIpPacket;
    private final BlockingQueue<IpPacket> queue = new LinkedBlockingQueue<>();

    private static final ExecutorService executor = Executors.newCachedThreadPool();;
    private  Future<?>  readFuture = null;
    private  Future<?>  writeFuture = null;

    private VpnService vpnService;

    public UdpConnect(VpnService vpnService, BytesParse bytesParse , IpPacket ipPacket) throws Exception {
        this.bytesParse = bytesParse;

        dstIpPacket = ipPacket;
        datagramSocket = new DatagramSocket();
        boolean protect = vpnService.protect(datagramSocket);
        datagramSocket.setSoTimeout(3000);
        this.vpnService = vpnService;
        readFuture = executor.submit(this::read);
        writeFuture = executor.submit(this::write);
    }

    @Override
    public void parseIpPacket(IpPacket ipPacket) {
        queue.offer(ipPacket);
    }

    private void read(){
        try{

            UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
            int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();

            while (!Thread.interrupted()){

                byte[] buffer = new byte[1024];
                DatagramPacket packet = new DatagramPacket(buffer, buffer.length);
                datagramSocket.receive(packet);

                //Push
                PacketTool.sendUdpPacket(this,buffer);

            }
        }catch (Exception e){
            close();
        }

    }

    private void write(){

        try {

            InetAddress dstAdder = dstIpPacket.getHeader().getDstAddr();
            UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
            int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();
            Random random = new Random();

            while (!Thread.interrupted()){
                IpPacket packet = queue.take();

                UdpPacket uPacket = (UdpPacket) packet.getPayload();
                byte[] rawData = uPacket.getPayload().getRawData();

                int port = uPacket.getHeader().getDstPort().valueAsInt();

                if(port == 53){
                    var dnsPacket = (DnsPacket)uPacket.getPayload();
                    DnsPacket.DnsHeader header = dnsPacket.getHeader();
                    short id = header.getId();

                    if (header.getQdCount() == 1) {
                        List<DnsQuestion> questions = header.getQuestions();
                        DnsQuestion dnsQuestion = questions.get(0);
                        DnsDomainName qName = dnsQuestion.getQName();
                        String domain = dnsQuestion.getQName().getName();
                        DnsResourceRecordType qType = dnsQuestion.getQType();

                        var myVpnService = (MyVpnService) vpnService;
                        String fakeIp = myVpnService.ipAndHostMap.inverse().get(domain);

                        if (fakeIp == null) {

                            for (int i = 0; i < 50; i++) {
                                String temp = random.nextInt(256) + "."
                                        + random.nextInt(256) + "."
                                        + random.nextInt(256) + "."
                                        + random.nextInt(256);

                                if (!myVpnService.ipAndHostMap.containsKey(temp)) {
                                    fakeIp = temp;

                                    myVpnService.ipAndHostMap.put(fakeIp, domain);

                                    break;
                                }
                            }

                            if (fakeIp == null) throw new RuntimeException("无法创建IP,稍后重试！！");

                        }

                        Inet4Address inetAddress = (Inet4Address)InetAddress.getByName(fakeIp); // 4 字节


                        DnsRDataA.Builder address = new DnsRDataA.Builder().address(inetAddress);
                        DnsRDataA build = address.build();

                        LengthBuilder<DnsResourceRecord> dnsResourceRecordLengthBuilder = new DnsResourceRecord.Builder()
                                .rdLength((short) 4)
                                .dataClass(DnsClass.IN).ttl(120)
                                .name(qName).dataType(qType)
                                .rData(build).correctLengthAtBuild(true);

                        DnsResourceRecord record = dnsResourceRecordLengthBuilder.build();

                        List<DnsResourceRecord> dnsResourceRecords = new ArrayList<>();
                        dnsResourceRecords.add(record);

                        DnsPacket.Builder dnsPacketBuild = new DnsPacket.Builder()
                                .qdCount((short) 1).anCount((short) 1)
                                .rCode(DnsRCode.NO_ERROR)
                                .opCode(DnsOpCode.QUERY)
                                .id(id).questions(questions).answers(dnsResourceRecords);

                        DnsPacket dPacket = dnsPacketBuild.build();

                        PacketTool.sendUdpPacket(this, dPacket.getRawData());
                    }
                }else{
                    DatagramPacket datagramPacket = new DatagramPacket(rawData, rawData.length, dstAdder, dstPort);
                    datagramSocket.send(datagramPacket);
                }

            }
        }catch (Exception e){
            close();

        }

    }


    private void close(){

        if(datagramSocket!=null){
            datagramSocket.close();
        }

        if(readFuture!=null){
            //发送中断指令
            readFuture.cancel(true);
        }

        if(writeFuture!=null){
            //发送中断指令
            writeFuture.cancel(true);
        }

        if(dstIpPacket!=null){
            bytesParse.removeKey(dstIpPacket);
        }

        if (dstIpPacket != null) {
            bytesParse.removeKey(dstIpPacket);
        }


    }

    @Override
    public AtomicInteger getSequenceNumber() {
        return null;
    }

    @Override
    public AtomicInteger getAckNumber() {
        return null;
    }

    @Override
    public void sendToVpn(byte[] packetData) {
        bytesParse.pushData(packetData);
    }

    @Override
    public InetSocketAddress getSrcAddress() {
        InetAddress dstAdder = dstIpPacket.getHeader().getDstAddr();
        UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
        int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();

        return new InetSocketAddress(dstAdder, dstPort);
    }

    @Override
    public InetSocketAddress getDstAddress() {
        InetAddress srcAdder = dstIpPacket.getHeader().getSrcAddr();
        UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
        int srcPort = udpPacket.getHeader().getSrcPort().valueAsInt();
        return new InetSocketAddress(srcAdder, srcPort);
    }


}
